import pandas as pd
import statsmodels.formula.api as smf
from sqlalchemy import create_engine
import matplotlib.pyplot as plt
import seaborn as sns

# 解决中文显示（仅保留系统大概率存在的字体）
plt.rcParams["font.family"] = ["SimHei", "Microsoft YaHei"]  # 优先系统常见中文字体
plt.rcParams["axes.unicode_minus"] = False  # 解决负号显示

# 数据库配置
db_config = {
    'host': 'localhost',
    'user': 'root',
    'password': 'sjk1234',
    'database': 'tushare',
    'port': 3306,
    'charset': 'utf8mb4'
}

# 连接数据库
engine = create_engine(
    f"mysql+pymysql://{db_config['user']}:{db_config['password']}@"
    f"{db_config['host']}:{db_config['port']}/{db_config['database']}?"
    f"charset={db_config['charset']}"
)
conn = engine.connect()

# 读取数据
df = pd.read_sql_query("""
    SELECT 
        d.*, 
        m.buy_lg_vol, m.sell_lg_vol, m.buy_elg_vol, m.sell_elg_vol, m.net_mf_vol,
        i.vol as i_vol, i.closes as i_closes 
    FROM date_1 d
    JOIN moneyflows m on d.ts_code = m.ts_code and d.trade_date = m.trade_date
    LEFT JOIN index_daily i on d.trade_date = i.trade_date and i.ts_code = '399001.SZ'
    WHERE d.trade_date BETWEEN '2023-01-01' AND '2023-12-31' 
      AND d.ts_code = '002229.SZ'
""", conn)

print("数据预览：")
print(df.head())

# 数据处理
numeric_fields = ['closes', 'i_closes', 'i_vol', 'vol', 'amount',
                 'buy_lg_vol', 'sell_lg_vol', 'buy_elg_vol', 'sell_elg_vol', 'net_mf_vol']
for field in numeric_fields:
    df[field] = pd.to_numeric(df[field], errors='coerce')

# 计算涨跌幅（避免自相关问题：删除zd_closes作为自变量）
df['zd_closes'] = round((df['closes'] - df['closes'].shift(1)) / df['closes'].shift(1), 4)
df['zs_closes'] = round((df['i_closes'] - df['i_closes'].shift(1)) / df['i_closes'].shift(1), 4)
df['zs_vol'] = round((df['i_vol'] - df['i_vol'].shift(1)) / df['i_vol'].shift(1), 4)
df = df.dropna(subset=['zd_closes', 'zs_closes', 'zs_vol']).reset_index(drop=True)

# 回归分析（关键修正：排除zd_closes自身作为自变量）
ex = ['id', 'ts_code', 'trade_date', 'the_date', 'opens', 'high', 'low',
      'closes', 'pre_closes', 'changes', 'pct_chg', 'zd_closes']  # 新增排除zd_closes
independent_vars = [col for col in df.select_dtypes(include=['number']).columns 
                   if col not in ex]

if independent_vars:
    formula = 'zd_closes ~ ' + ' + '.join(independent_vars)
    res = smf.ols(formula, data=df).fit()
    print("\n回归关键指标：")
    print(f"R方: {res.rsquared:.4f}")
    print(f"调整后R方: {res.rsquared_adj:.4f}")
    print(f"F统计量: {res.fvalue:.2f} (P值: {res.f_pvalue:.4f})")
    print("\n核心系数（P值<0.1的变量）：")
    # 只显示P值显著的变量
    sig_params = res.pvalues[res.pvalues < 0.1].index
    if len(sig_params) > 0:
        print(res.params[sig_params].round(4))
    else:
        print("无显著相关的变量（P值<0.1）")

# 可视化（确保中文正常）
plt.figure(figsize=(10, 6))
sns.scatterplot(x='zd_closes', y='zs_closes', data=df)
plt.title('个股涨跌幅 vs 指数涨跌幅')
plt.xlabel('个股涨跌幅')
plt.ylabel('指数涨跌幅')
plt.tight_layout()  # 调整布局避免截断
plt.show()

# 相关性矩阵
plt.figure(figsize=(12, 8))
corr = df[['vol', 'amount', 'buy_lg_vol', 'sell_lg_vol', 
           'buy_elg_vol', 'sell_elg_vol', 'net_mf_vol', 'zs_vol']].corr()
sns.heatmap(corr, annot=True, cmap='coolwarm', fmt='.2f', linewidths=.5)
plt.title('变量相关性矩阵')
plt.tight_layout()
plt.show()

conn.close()
